-module (osmos_tree).

-export ([ new/1,
	   search/2,
	   insert/4,
	   update/4,
	   total_entries/1,
	   iter_begin/1,
	   iter_end/1,
	   iter_lower_bound/2,
	   iter_next/1,
	   iter_key/1,
	   iter_compare_key/1,
	   iter_value/1,
	   iter_less/2 ]).

-include ("osmos_compare.hrl").

-record (tree, { compare, tree }).
-record (iter, { key, key_cmp, value, compare, iter }).

-define (ITER_END, '$end_of_table').

new (Cmp) ->
  #tree { compare = Cmp, tree = my_gb_trees:empty () }.

search (#tree { compare = Cmp, tree = Tree }, KeyCmp) ->
  case my_gb_trees:lookup (Cmp#osmos_compare.less, KeyCmp, Tree) of
    { value, { K, V } } ->
      { ok, { K, V } };
    none ->
      not_found
  end.

insert (T = #tree { compare = Cmp, tree = Tree }, Key, KeyCmp, Value) ->
  NewTree = my_gb_trees:insert (Cmp#osmos_compare.less, KeyCmp, { Key, Value }, Tree),
  T#tree { tree = NewTree }.

update (T = #tree { compare = Cmp, tree = Tree }, Key, KeyCmp, Value) ->
  NewTree = my_gb_trees:update (Cmp#osmos_compare.less, KeyCmp, { Key, Value }, Tree),
  T#tree { tree = NewTree }.

total_entries (#tree { tree = Tree }) ->
  my_gb_trees:size (Tree).

iter_begin (#tree { compare = Cmp, tree = Tree }) ->
  case my_gb_trees:next (my_gb_trees:iterator (Tree)) of
    { KCmp, { K, V }, It } ->
      #iter { key = K, key_cmp = KCmp, value = V, compare = Cmp, iter = It };
    none ->
      ?ITER_END
  end.

iter_end (_Tree) ->
  ?ITER_END.

iter_lower_bound (Tree, Less) when is_function (Less) ->
  % XXX: inefficient
  iter_lower_bound_scan (iter_begin (Tree), Less);

iter_lower_bound (Tree, KeyCmp) ->
  iter_lower_bound_scan (iter_begin (Tree), fun (K) -> K < KeyCmp end).

iter_lower_bound_scan (?ITER_END, _Less) ->
  ?ITER_END;
iter_lower_bound_scan (It, Less) ->
% XXX: using iter_compare_key here is kind of weird
  case Less (iter_compare_key (It)) of
    false -> It;
    true  -> iter_lower_bound_scan (iter_next (It), Less)
  end.

iter_next (#iter { compare = Cmp, iter = It }) ->
  case my_gb_trees:next (It) of
    { KCmp, { K, V }, NewIt } ->
      #iter { key = K, key_cmp = KCmp, value = V, compare = Cmp, iter = NewIt };
    none ->
      ?ITER_END
  end.

iter_key (#iter { key = K }) ->
  K.
iter_compare_key (#iter { key_cmp = K }) ->
  K.
iter_value (#iter { value = V }) ->
  V.

iter_less (?ITER_END, ?ITER_END) ->
  false;
iter_less (?ITER_END, B) when is_record (B, iter) ->
  false;
iter_less (A, ?ITER_END) when is_record (A, iter) ->
  true;
iter_less (#iter { key_cmp = A, compare = Cmp },
	   #iter { key_cmp = B, compare = Cmp }) ->
  ?osmos_compare_less (Cmp, A, B).
